import os
import cv2 as cv
from zipfile import ZipFile
from skimage import io
import numpy as np
import random
import logging
import yaml

logging.basicConfig(level=logging.INFO)

def random_transform(img, scale_min=0.5, scale_max=1.5):
    h, w = img.shape[:2]
    scale_factor = random.uniform(scale_min, scale_max)
    new_h, new_w = int(h * scale_factor), int(w * scale_factor)
    
    # 缩放图像
    scaled_img = cv.resize(img, (new_w, new_h), interpolation=cv.INTER_LINEAR)
    
    # 中心裁剪或填充
    if scale_factor >= 1:
        y_start = (new_h - h) // 2
        x_start = (new_w - w) // 2
        cropped_img = scaled_img[y_start:y_start+h, x_start:x_start+w]
    else:
        cropped_img = np.zeros_like(img)
        y_start = (h - new_h) // 2
        x_start = (w - new_w) // 2
        cropped_img[y_start:y_start+new_h, x_start:x_start+new_w] = scaled_img
    
    # 随机旋转（-45° ~ 45°）
    angle = random.uniform(-45, 45)
    M = cv.getRotationMatrix2D((w//2, h//2), angle, 1)
    rotated_img = cv.warpAffine(cropped_img, M, (w, h), borderMode=cv.BORDER_REPLICATE)
    
    return rotated_img

def process_images_from_zip(
    zip_path, 
    scale_min=0.5, 
    scale_max=1.5, 
    num_images=10, 
    output_type='RGB',
    progress_callback=None,
    thread_event=None
):
    out_path = 'out/'
    os.makedirs(out_path, exist_ok=True)
    logging.info(f"开始处理 ZIP 文件: {zip_path}")
    
    with ZipFile(zip_path, 'r') as zipped:
        image_files = [
            name for name in zipped.namelist()
            if name.lower().endswith(('.png', '.jpg', '.jpeg'))  # 仅处理JPEG/PNG
        ]
        total = len(image_files)
        
        for idx, name in enumerate(image_files):
            if thread_event and thread_event.is_set():
                logging.info("用户取消处理")
                return
            
            # 更新进度
            if progress_callback:
                progress_callback(int((idx + 1) / total * 100))
            
            with zipped.open(name) as file:
                # 读取图像
                img_bytes = file.read()
                img = cv.imdecode(np.frombuffer(img_bytes, np.uint8), -1)
                
                # 处理输出类型
                if output_type == 'Grayscale':
                    img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
                elif output_type == 'Binary':
                    img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
                    _, img = cv.threshold(img, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)
                
                # 调整尺寸
                img = cv.resize(img, (224, 224))
                
                # 保存原始处理结果
                base_name = os.path.splitext(os.path.basename(name))[0]
                save_path = os.path.join(out_path, f"{base_name}_original.jpg")
                io.imsave(save_path, img)
                
                # 生成变换图像
                for i in range(num_images - 1):
                    transformed = random_transform(img, scale_min, scale_max)
                    trans_path = os.path.join(out_path, f"{base_name}_transformed_{i}.jpg")
                    io.imsave(trans_path, transformed)
    
    logging.info(f"处理完成，结果保存至: {out_path}")

def load_config(config_path='config.yaml'):
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    except FileNotFoundError:
        logging.warning("配置文件未找到，使用默认值")
        return {
            "scale_min": 0.5,
            "scale_max": 1.5,
            "num_images": 10,
            "output_type": "RGB"
        }
    except yaml.YAMLError as e:
        logging.error(f"配置文件解析错误: {e}")
        return {}